import numpy as np
import pyaudio
import time
import soundfile as sf
from voxcpm import VoxCPM # 导入 VoxCPM 模型类
from typing import Iterator, Tuple

# 注意：VoxCPM 默认采样率为 16000 Hz。
VOXCPM_SAMPLE_RATE = 16000

class VoxCPMStreamer:
    """
    VoxCPM 流式文本到语音 (TTS) 合成器。
    使用 model.generate_streaming 实时获取音频片段，并使用 PyAudio 连续播放。
    """
    def __init__(self, model_id="openbmb/VoxCPM-0.5B", sample_rate=VOXCPM_SAMPLE_RATE):
        self.sample_rate = sample_rate
        # 初始化 PyAudio
        self.p = pyaudio.PyAudio()
        
        print(f"正在加载 VoxCPM 模型: {model_id}...")
        # 1. 加载 VoxCPM 模型
        try:
            self.model = VoxCPM.from_pretrained(model_id)
            print("✅ 模型加载成功。")
        except Exception as e:
            print(f"致命错误: VoxCPM 模型加载失败: {e}")
            raise

    def _voxcpm_generator(self, text: str, **kwargs) -> Iterator[np.ndarray]:
        """
        封装 VoxCPM 的 generate_streaming 方法，直接返回音频块的生成器。
        """
        # 2. 调用 VoxCPM 的流式生成方法
        # 传递额外的参数，例如降噪、推理步数等
        return self.model.generate_streaming(text=text, **kwargs)

    def stream_synthesis(self, text: str, **kwargs):
        """
        使用 VoxCPM 进行流式合成和连续播放。
        """
        
        generator = self._voxcpm_generator(text, **kwargs)
        
        print("开始 VoxCPM 流式合成...")
        
        # 【优化点 1】: 在循环外只打开一次 PyAudio 流，解决片段间的卡顿问题
        stream = self.p.open(
            format=pyaudio.paInt16, # 适配 VoxCPM 的 16-bit PCM 输出格式
            channels=1,
            rate=self.sample_rate,
            output=True,
            frames_per_buffer=4096,
        )
        
        try:
            for i, chunk in enumerate(generator):
                # 这里的 chunk 是一个 numpy 数组 (float32)，即模型生成的音频片段
                
                # 3. 将 float32 转换为 Int16 格式 (VoxCPM默认输出)，并进行归一化/缩放
                audio_int16 = (chunk * 32767).astype(np.int16) 
                
                # 写入音频数据块
                stream.write(audio_int16.tobytes())
                
                print(f"\n🔊 播放片段 {i+1}: 已接收 {len(chunk)} 个样本。")
            
        except Exception as e:
            print(f"错误: 音频播放失败: {e}")
        finally:
            # 【优化点 1】: 在循环结束后只关闭一次 PyAudio 流
            if 'stream' in locals() and stream.is_active():
                stream.stop_stream()
                stream.close()
        
        print("\n✅ VoxCPM 流式合成完成！")
    
    def close(self):
        """释放 PyAudio 资源"""
        if hasattr(self, 'p'):
            self.p.terminate()

# 使用示例
if __name__ == "__main__":
    # 示例文本
    test_text = """欢迎使用 VoxCPM 流式合成。
    这个模型是一个创新的端到端 TTS 模型，旨在生成具有高度表现力的语音。
    流式模式可以实现低延迟的实时语音输出，提供更佳的用户体验。"""
    prompt_wav_path = r'TTS/我此番也是受神子之邀，体验一下市井游乐的氛围，和各位并无二致。.wav'
    prompt_text = '我此番也是受神子之邀，体验一下市井游乐的氛围，和各位并无二致。'
    # 初始化 VoxCPM 流式合成器
    streamer = VoxCPMStreamer()
    try:
        # 【优化点 2】: 降低 inference_timesteps 到 6 或 7，显著提高流式速度。
        # 较低的值会损失一些质量，但能大大改善实时性。
        print("\n--- 优化设置：使用 inference_timesteps=7 以提高速度 ---")
        streamer.stream_synthesis(
            test_text, 
            prompt_wav_path=prompt_wav_path,
            prompt_text=prompt_text,    
            denoise=True,                 # 启用降噪工具 (推荐)
            inference_timesteps=7,        # 速度优化，推荐 7-8
            cfg_value=2.0                 
        )
        
    except Exception as e:
        print(f"程序运行错误: {e}")
    finally:
        # 确保资源被正确释放
        streamer.close()
